import numpy as np
from pykalman import KalmanFilter


class VehicleKalmanFilter:
    """
    车辆运动卡尔曼滤波器，优化用于车辆加速度和角速度数据处理
    """

    def __init__(self, accel_range=2.0, gyro_range=2000.0):
        """初始化车辆运动卡尔曼滤波器"""
        # 传感器参数
        self.accel_range = accel_range  # g
        self.gyro_range = gyro_range  # dps

        # ICM42688噪声特性
        self.gyro_noise = 2.8e-3  # 2.8 mdps/√Hz
        self.accel_noise = 70e-6  # 70 μg/√Hz

        # 状态变量: [accel_x, accel_y, accel_z, gyro_x, gyro_y, gyro_z]
        self.state_dim = 6

        # 观测变量: [accel_x, accel_y, accel_z, gyro_x, gyro_y, gyro_z]
        self.obs_dim = 6

        # 初始化滤波器
        self.kf = self._init_kalman_filter()
        self.filter_state = None
        self.filter_covariance = None

        # 车辆特定参数
        self.last_gravity = np.array([0, 0, 9.8])  # 初始重力方向
        self.vehicle_motion_state = "STATIONARY"  # 初始状态
        self.stationary_threshold = 0.05  # 静止状态阈值(g)
        self.last_update_time = None

    def _init_kalman_filter(self):
        """初始化卡尔曼滤波器参数"""
        # 状态转移矩阵 - 假设传感器数据在短时间内有一定相关性
        transition_matrix = np.eye(self.state_dim) * 0.98  # 轻微衰减

        # 观测矩阵 - 直接观测状态
        observation_matrix = np.eye(self.obs_dim)

        # 状态转移噪声协方差
        transition_covariance = np.diag([
            self.accel_noise ** 2 * 500,  # accel_x 过程噪声
            self.accel_noise ** 2 * 500,  # accel_y 过程噪声
            self.accel_noise ** 2 * 500,  # accel_z 过程噪声
            self.gyro_noise ** 2 * 500,  # gyro_x 过程噪声
            self.gyro_noise ** 2 * 500,  # gyro_y 过程噪声
            self.gyro_noise ** 2 * 500  # gyro_z 过程噪声
        ])

        # 观测噪声协方差
        observation_covariance = np.diag([
            self.accel_noise ** 2 * 200,  # accel_x 测量噪声
            self.accel_noise ** 2 * 200,  # accel_y 测量噪声
            self.accel_noise ** 2 * 200,  # accel_z 测量噪声
            self.gyro_noise ** 2 * 200,  # gyro_x 测量噪声
            self.gyro_noise ** 2 * 200,  # gyro_y 测量噪声
            self.gyro_noise ** 2 * 200  # gyro_z 测量噪声
        ])

        # 初始状态协方差
        initial_covariance = np.eye(self.state_dim) * 0.1

        # 初始化卡尔曼滤波器
        kf = KalmanFilter(
            transition_matrices=transition_matrix,
            observation_matrices=observation_matrix,
            transition_covariance=transition_covariance,
            observation_covariance=observation_covariance,
            initial_state_covariance=initial_covariance
        )

        return kf

    def update(self, acceleration, gyroscope, dt):
        """
        使用新的传感器数据更新卡尔曼滤波器

        参数:
            acceleration: 加速度数据 [x, y, z] (g)
            gyroscope: 角速度数据 [x, y, z] (dps)
            dt: 时间间隔 (秒)

        返回:
            filtered_state: 滤波后的状态 [accel_x, accel_y, accel_z, gyro_x, gyro_y, gyro_z]
        """
        # 构建观测向量
        observation = np.concatenate([acceleration, gyroscope])

        # 自适应调整滤波器参数
        self._adapt_filter_parameters(acceleration, gyroscope, dt)

        # 初始化状态
        if self.filter_state is None:
            self.filter_state = observation
            self.filter_covariance = self.kf.initial_state_covariance

        # 执行卡尔曼滤波预测和更新步骤
        predicted_state, predicted_covariance = self.kf.filter_update(
            filtered_state_mean=self.filter_state,
            filtered_state_covariance=self.filter_covariance,
            observation=observation
        )

        # 更新内部状态
        self.filter_state = predicted_state
        self.filter_covariance = predicted_covariance

        # 更新重力方向估计
        self._update_gravity_estimate()

        return self.filter_state

    def _adapt_filter_parameters(self, acceleration, gyroscope, dt):
        """根据车辆运动状态自适应调整滤波器参数"""
        # 计算加速度大小
        accel_magnitude = np.linalg.norm(acceleration)

        # 检测静止状态
        if accel_magnitude < self.stationary_threshold and np.linalg.norm(gyroscope) < 5.0:
            self.vehicle_motion_state = "STATIONARY"
            # 静止状态下增加观测权重
            self.kf.observation_covariance = np.diag([
                self.accel_noise ** 2 * 100,  # 更高的观测权重
                self.accel_noise ** 2 * 100,
                self.accel_noise ** 2 * 100,
                self.gyro_noise ** 2 * 100,
                self.gyro_noise ** 2 * 100,
                self.gyro_noise ** 2 * 100
            ])
        else:
            self.vehicle_motion_state = "MOVING"
            # 运动状态下增加过程噪声
            self.kf.transition_covariance = np.diag([
                self.accel_noise ** 2 * 800,  # 更高的过程噪声
                self.accel_noise ** 2 * 800,
                self.accel_noise ** 2 * 800,
                self.gyro_noise ** 2 * 800,
                self.gyro_noise ** 2 * 800,
                self.gyro_noise ** 2 * 800
            ])

    def _update_gravity_estimate(self):
        """更新重力方向估计"""
        if self.vehicle_motion_state == "STATIONARY":
            # 静止状态下，加速度主要是重力
            self.last_gravity = self.filter_state[:3]
            # 归一化
            norm = np.linalg.norm(self.last_gravity)
            if norm > 0:
                self.last_gravity = self.last_gravity / norm

    def get_filtered_acceleration(self, raw_acceleration):
        """
        从滤波后的状态中获取优化的加速度值，去除重力分量

        参数:
            raw_acceleration: 原始加速度数据 [x, y, z] (g)

        返回:
            filtered_acceleration: 滤波后的加速度 [x, y, z] (g)
        """
        # 获取滤波后的加速度
        filtered_accel = self.filter_state[:3]

        # 估计重力分量
        gravity_component = np.dot(filtered_accel, self.last_gravity) * self.last_gravity

        # 减去重力分量得到纯运动加速度
        motion_acceleration = filtered_accel

        return motion_acceleration

    def get_filtered_gyroscope(self, convert_to_rad=False):
        """
        获取滤波后的角速度数据

        参数:
            convert_to_rad: 是否将结果从度/秒转换为弧度/秒

        返回:
            filtered_gyro: 滤波后的角速度 [x, y, z] (单位由convert_to_rad参数决定)
        """
        # 提取滤波后的角速度
        filtered_gyro = self.filter_state[3:]

        # 应用偏置校正（如果有）
        if hasattr(self, 'gyro_bias') and self.gyro_bias is not None:
            filtered_gyro = filtered_gyro - self.gyro_bias

        # 单位转换（如果需要）
        if convert_to_rad:
            filtered_gyro = np.radians(filtered_gyro)

        return filtered_gyro

    def get_vehicle_motion_metrics(self):
        """获取车辆运动指标"""
        # 提取滤波后的加速度和角速度
        accel = self.filter_state[:3]
        gyro = self.filter_state[3:]

        # 计算加速度大小（去除重力）
        motion_accel = self.get_filtered_acceleration(accel)
        accel_magnitude = np.linalg.norm(motion_accel)

        # 计算转向角速度大小
        turning_rate = np.abs(gyro[2])  # z轴角速度（偏航）

        return {
            "acceleration_magnitude": accel_magnitude,
            "turning_rate": turning_rate,
            "motion_state": self.vehicle_motion_state
        }
